#!/usr/bin/env python3
"""
AudioReceiver - 音频接收和处理模块
支持自动采样率检测和重采样，确保音频同步
"""

import asyncio
import time
import numpy as np
import av
from aiortc_custom.logger_config import log_info, log_debug, log_warning, log_error, log_success, log_failure, log_progress, log_connection, log_media, log_stats
from FFmpegAudioPlayer import FFmpegAudioPlayer
from av import AudioFrame

class AudioReceiver:
    """音频接收器 - 支持采样率检测、重采样和缓冲管理"""

    def __init__(self):
        self.frame_count = 0
        self.last_stats_time = 0
        self.stats_interval = 1.0
        self.pc = None
        self.audio_player = None
        self.player_start_time = None

        # 重采样相关
        self.resampler = None
        self.source_sample_rate = None
        self.source_channels = None
        self.target_sample_rate = 48000  # FFplay使用48kHz
        self.target_channels = 2  # 统一为双声道
        self.resampling_needed = False

        # 缓冲统计
        self.resampled_frames = 0
        self.dropped_frames = 0

    def set_peer_connection(self, pc):
        """设置对等连接"""
        self.pc = pc

    async def handle_track(self, track):
        """处理接收的远程音频轨道"""
        log_success("开始处理接收的远程音频轨道")

        # 初始化FFmpeg音频播放器
        try:
            # 获取音频参数
            sample_rate, channels = await self._get_audio_params(track)

            log_info(f"检测到音频参数: {sample_rate}Hz, {channels}声道")

            # 初始化重采样器（如果需要）
            self._init_resampler(sample_rate, channels)

            # 创建播放器
            self.audio_player = FFmpegAudioPlayer(buffer_size=10)  # 增加缓冲以处理重采样
            success = self.audio_player.start(
                sample_rate=self.target_sample_rate,
                channels=self.target_channels
            )

            if not success:
                log_failure("FFmpeg音频播放器启动失败")
                return

            self.player_start_time = time.time()

            if self.resampling_needed:
                log_success(
                    f"FFmpeg音频播放器启动成功: {sample_rate}Hz -> {self.target_sample_rate}Hz, "
                    f"{channels} -> {self.target_channels}声道（需要重采样）"
                )
            else:
                log_success(
                    f"FFmpeg音频播放器启动成功: {sample_rate}Hz, {channels}声道"
                )

        except Exception as e:
            log_failure(f"初始化音频播放器失败: {e}")
            import traceback
            traceback.print_exc()
            return

        log_success("开始接收远程音频帧...")

        consecutive_errors = 0
        max_consecutive_errors = 10

        while True:
            try:
                # 检查轨道状态
                if track.readyState == "ended":
                    log_info("远程音频轨道已结束")
                    break

                # 接收音频帧
                frame = await track.recv()
                if frame is None:
                    log_info("音频轨道结束，停止接收")
                    break

                # 处理音频帧
                await self._process_audio_frame(frame)

                # 重置错误计数
                consecutive_errors = 0

            except Exception as e:
                consecutive_errors += 1
                log_debug(f"音频处理错误 ({consecutive_errors}/{max_consecutive_errors}): {e}")

                if consecutive_errors >= max_consecutive_errors:
                    log_error("音频处理连续错误过多，停止处理")
                    break

                # 短暂等待后继续
                await asyncio.sleep(0.01)

        # 清理资源
        await self._cleanup()
        log_success("音频接收处理完成")

    async def _get_audio_params(self, track):
        """获取音频参数"""
        try:
            # 尝试从轨道获取参数
            sample_rate = getattr(track, 'sample_rate', None)
            if not sample_rate or sample_rate <= 0:
                log_debug("轨道不提供采样率，使用默认值48000")
                sample_rate = 48000

            # 尝试获取声道数
            channels = 1  # 默认单声道
            if hasattr(track, 'layout') and track.layout:
                try:
                    if hasattr(track.layout, 'channels'):
                        layout_channels = track.layout.channels
                        if isinstance(layout_channels, (list, tuple)):
                            if layout_channels and int(layout_channels[0]) > 0:
                                channels = int(layout_channels[0])
                        else:
                            ch_value = int(layout_channels)
                            if ch_value > 0:
                                channels = ch_value
                except (ValueError, TypeError) as e:
                    log_debug(f"无法解析layout.channels: {e}")
                    channels = 1

            elif hasattr(track, 'channels'):
                try:
                    track_channels = track.channels
                    if isinstance(track_channels, (list, tuple)):
                        if track_channels and int(track_channels[0]) > 0:
                            channels = int(track_channels[0])
                    else:
                        ch_value = int(track_channels)
                        if ch_value > 0:
                            channels = ch_value
                except (ValueError, TypeError) as e:
                    log_debug(f"无法解析track.channels: {e}")
                    channels = 1

            return sample_rate, channels

        except Exception as e:
            log_debug(f"获取音频参数失败: {e}")
            return 48000, 1

    def _init_resampler(self, sample_rate, channels):
        """初始化重采样器"""
        self.source_sample_rate = sample_rate
        self.source_channels = channels

        # 检查是否需要重采样
        if sample_rate != self.target_sample_rate or channels != self.target_channels:
            self.resampling_needed = True

            try:
                log_info(f"初始化重采样器: {sample_rate}Hz/{channels}ch -> {self.target_sample_rate}Hz/{self.target_channels}ch")

                # 创建重采样器
                self.resampler = av.AudioResampler(
                    format='s16',
                    layout='stereo' if self.target_channels == 2 else 'mono',
                )
                log_success("重采样器初始化成功")

            except Exception as e:
                log_error(f"重采样器初始化失败: {e}")
                self.resampler = None
                self.resampling_needed = False
        else:
            self.resampling_needed = False
            log_info("源音频参数与目标参数匹配，无需重采样")

    async def _process_audio_frame(self, frame):
        """处理单个音频帧"""
        try:
            if not isinstance(frame, AudioFrame):
                return

            # 如果音频播放器未初始化，自动初始化
            if self.audio_player is None:
                sample_rate = getattr(frame, 'sample_rate', 48000)
                channels = 1

                if hasattr(frame, 'layout') and frame.layout:
                    try:
                        if hasattr(frame.layout, 'channels'):
                            layout_channels = frame.layout.channels
                            if isinstance(layout_channels, (list, tuple)):
                                if layout_channels and int(layout_channels[0]) > 0:
                                    channels = int(layout_channels[0])
                            else:
                                ch_value = int(layout_channels)
                                if ch_value > 0:
                                    channels = ch_value
                    except (ValueError, TypeError):
                        channels = 1

                # 初始化重采样器
                self._init_resampler(sample_rate, channels)

                # 创建播放器
                from FFmpegAudioPlayer import FFmpegAudioPlayer
                self.audio_player = FFmpegAudioPlayer(buffer_size=10)
                success = self.audio_player.start(
                    sample_rate=self.target_sample_rate,
                    channels=self.target_channels
                )

                if success:
                    self.player_start_time = time.time()
                    if self.resampling_needed:
                        log_success(
                            f"FFmpeg音频播放器自动启动: {sample_rate}Hz -> {self.target_sample_rate}Hz "
                            f"（需要重采样）"
                        )
                    else:
                        log_success(f"FFmpeg音频播放器自动启动: {sample_rate}Hz, {channels}声道")
                else:
                    log_failure("FFmpeg音频播放器自动启动失败")
                    return

            # 转换音频数据
            pcm = self._convert_audio_frame(frame)

            # 检查音频数据是否有效
            if pcm is None or pcm.size == 0:
                return

            # 进行重采样（如果需要）
            if self.resampling_needed and self.resampler:
                try:
                    # 创建av AudioFrame以便进行重采样
                    # 重要：av库对shape的要求：mono是(samples,)，stereo是(samples, 2)
                    if self.source_channels == 1:
                        # 单声道：保持1D shape
                        audio_data = pcm.reshape(-1)
                    else:
                        # 多声道：保持2D shape
                        audio_data = pcm.reshape(-1, self.source_channels)

                    av_frame = av.AudioFrame.from_ndarray(
                        audio_data,
                        format='s16',
                        layout='mono' if self.source_channels == 1 else 'stereo'
                    )
                    av_frame.sample_rate = self.source_sample_rate

                    # 进行重采样
                    resampled_frames = self.resampler.resample(av_frame)

                    if resampled_frames:
                        for resampled_frame in resampled_frames:
                            resampled_pcm = resampled_frame.to_ndarray()

                            # 确保内存布局连续
                            if not resampled_pcm.flags['C_CONTIGUOUS']:
                                resampled_pcm = np.ascontiguousarray(resampled_pcm, dtype='int16')

                            # 发送到FFmpeg播放器
                            if self.audio_player is not None:
                                self.audio_player.write_frame(resampled_pcm)
                                self.resampled_frames += 1

                except Exception as e:
                    log_warning(f"重采样失败，尝试直接使用原始数据: {e}")
                    # 如果重采样失败，尝试直接使用数据
                    if self.audio_player is not None:
                        self.audio_player.write_frame(pcm)
            else:
                # 不需要重采样，直接使用数据
                # 确保内存布局连续
                if not pcm.flags['C_CONTIGUOUS']:
                    pcm = np.ascontiguousarray(pcm, dtype='int16')

                # 发送到FFmpeg播放器
                if self.audio_player is not None:
                    self.audio_player.write_frame(pcm)

            self.frame_count += 1

            # 定期输出统计信息
            if self.frame_count % 100 == 0:
                await self._update_stats()

        except Exception as e:
            log_warning(f"音频帧处理错误: {e}")

    def _convert_audio_frame(self, frame):
        """转换音频帧为numpy数组"""
        try:
            pcm = frame.to_ndarray()

            # 确保是 int16 格式
            if pcm.dtype != np.int16:
                if pcm.dtype in [np.float32, np.float64]:
                    pcm = (pcm * 32767).clip(-32768, 32767).astype(np.int16)
                elif pcm.dtype in [np.int8, np.uint8]:
                    pcm = (pcm.astype(np.int16) * 256).clip(-32768, 32767)
                elif pcm.dtype in [np.int32, np.int64]:
                    pcm = (pcm / (2**(pcm.dtype.itemsize * 8 - 16))).astype(np.int16)
                else:
                    pcm = pcm.astype(np.int16)

            return pcm

        except Exception as e:
            log_warning(f"音频帧转换失败: {e}")
            # 返回静音数据
            return np.zeros(1024, dtype=np.int16)

    async def _update_stats(self):
        """更新统计信息"""
        try:
            if self.audio_player:
                stats = self.audio_player.get_stats()
                if self.resampling_needed:
                    log_stats(
                        f"音频统计: 已处理{self.frame_count}帧, 重采样{self.resampled_frames}帧, "
                        f"丢弃{stats.get('dropped_frames', 0)}帧, "
                        f"队列{stats.get('queue_size', 0)}, "
                        f"丢包率{stats.get('drop_rate', 0):.2f}%"
                    )
                else:
                    log_stats(
                        f"音频统计: 已处理{self.frame_count}帧, "
                        f"丢弃{stats.get('dropped_frames', 0)}帧, "
                        f"队列{stats.get('queue_size', 0)}, "
                        f"丢包率{stats.get('drop_rate', 0):.2f}%"
                    )
        except Exception as e:
            log_debug(f"更新统计失败: {e}")

    async def _cleanup(self):
        """清理资源"""
        try:
            if self.audio_player is not None:
                log_info("正在停止音频播放器...")
                self.audio_player.stop()
                self.audio_player = None

            # 清理重采样器
            self.resampler = None

            log_success("音频播放器已清理")
        except Exception as e:
            log_info(f"清理音频播放器时出错: {e}")

    def get_stats(self):
        """获取统计信息"""
        if self.audio_player is not None:
            stats = self.audio_player.get_stats()
            return {
                'frame_count': self.frame_count,
                'resampled_frames': self.resampled_frames,
                'resampling_needed': self.resampling_needed,
                'source_rate': self.source_sample_rate,
                'source_channels': self.source_channels,
                'target_rate': self.target_sample_rate,
                'target_channels': self.target_channels,
                'is_playing': stats.get('is_running', False),
                'player_active': self.audio_player is not None,
                'dropped_frames': stats.get('dropped_frames', 0),
                'queue_size': stats.get('queue_size', 0),
                'drop_rate': stats.get('drop_rate', 0)
            }
        else:
            return {
                'frame_count': self.frame_count,
                'resampled_frames': self.resampled_frames,
                'resampling_needed': self.resampling_needed,
                'source_rate': self.source_sample_rate,
                'source_channels': self.source_channels,
                'target_rate': self.target_sample_rate,
                'target_channels': self.target_channels,
                'is_playing': False,
                'player_active': False,
                'dropped_frames': 0,
                'queue_size': 0,
                'drop_rate': 0
            }
